from torch.utils.data import Dataset,DataLoader
import cv2
import os
import numpy as np
import torch
import torch.nn as nn
from pathlib import Path
# import skimage.io as io
# from skimage.transform import resize
from process_data import concat_data,get_data_list,get_binary_list_stage1,get_binary_list_stage2,normalize_data
from data_aug import aug_data,random_rotate_and_crop
import random


def mixup_criterion(criterion, pred, y_a, y_b, lam):
    return lam * criterion(pred, y_a) + (1 - lam) * criterion(pred, y_b)

def mixup_data(x, y, alpha=1., use_cuda=True):
    '''Returns mixed inputs, pairs of targets, and lambda'''
    if alpha > 0:
        lam = np.random.beta(alpha, alpha)
    else:
        lam = 1

    batch_size = x.size()[0]
    if use_cuda:
        index = torch.randperm(batch_size).cuda()
    else:
        index = torch.randperm(batch_size)

    mixed_x = lam * x + (1 - lam) * x[index, :]
    y_a, y_b = y, y[index]
    return mixed_x, y_a, y_b, lam

def make_category_list(data_list):
    ## 将data list转成list中包含4个list，分别对应四个类别。每个list中包含该类别的图片关键帧路径。
    cate_list_0 = [p for p in data_list if p[1] == 0]
    cate_list_1 = [p for p in data_list if p[1] == 1]
    cate_list_2 = [p for p in data_list if p[1] == 2]
    cate_list_3 = [p for p in data_list if p[1] == 3]
    return [cate_list_0,cate_list_1,cate_list_2,cate_list_3]

class Mydataset(Dataset):
    def __init__(self,data_list,train = True,image_size = None,concat = False,normalize = False,shuffle_concat = False):
        ## data_list中保存每个数据关键帧的路径和对应的label
        self.data_list = data_list
        self.shuffle_concat = shuffle_concat
        if shuffle_concat:
            self.category_list = make_category_list(data_list)
        self.image_size = image_size
        ## 是否将一个序列下的3-5张图片进行concat操作。组合到一起作为一张图片进行输入。
        ## concat为False表示不进行concat操作，=0表示从垂直方向concat，=1表示从水平方向concat
        self.concat = concat
        self.normalize = normalize
        self.train = train
    def process_data(self,index):
        ## 是否concat来判断是否将图片进行concat操作。进行concat操作需要传入的list是key_frame = True,
        ## 这样才可以取出某个文件夹下所有数据。
        if self.concat is not False:
            label = self.data_list[index][1]
            ## 直接获取该路径的上一层目录路径
            data_path = os.path.dirname(self.data_list[index][0])
            ## 从路径中获取关键帧
            key_frame = int(self.data_list[index][0].replace('\\','/').split('/')[-1].replace('.jpg',''))
            data_list = []
            if self.train:
                name_list = os.listdir(data_path)
                if self.shuffle_concat:
                    # 在原始的路径列表中取出一个
                    name_list.remove(random.choice(name_list))
                    ## 然后在同等类别的list中随机选择一个加入.注意选择对应label中的第一个路径元素。
                    name_list.append(random.choice(self.category_list[label])[0])
                random.shuffle(name_list)
                # print(name_list,'*'*20)
            else:
                name_list = os.listdir(data_path)
            for name in name_list:
                if '/' in name:
                    image_dir = name
                else:
                    image_dir = os.path.join(data_path,name)
                image = cv2.resize(cv2.imread(image_dir), self.image_size)
                # image = resize(io.imread(image_dir),self.image_size[::-1])
                if self.train:
                    image = aug_data(image, image_size=self.image_size, p=1.)
                data_list.append(image)
            image = concat_data(data_list,key_frame,down_sample = 1,concat=self.concat)
            if self.normalize:
                image = normalize_data(image)
            else:
                image = torch.from_numpy(np.transpose(image, (2, 0, 1))).float()

        else:
            image_dir = self.data_list[index][0]
            image = cv2.resize(cv2.imread(image_dir), self.image_size)
            # image = resize(io.imread(image_dir),self.image_size[::-1])
            label = self.data_list[index][1]
            label = torch.tensor(label)
            image = aug_data(image, image_size=self.image_size, p=1.)
            if self.normalize:
                image = normalize_data(image)
            else:
                image = torch.from_numpy(np.transpose(image, (2, 0, 1))).float()
        return image,label

    def __getitem__(self, index):
        image,label = self.process_data(index=index)
        return image,label


    def __len__(self):
        return len(self.data_list)



if __name__ == '__main__':
    data_list_train,data_list_val = get_binary_list_stage1(only_key_frame=True)
    dataset = Mydataset(data_list=data_list_train,image_size=(960,512),concat=True)
    data_loader = DataLoader(dataset=dataset,batch_size=2,shuffle=True)
    for batch_image,batch_label in data_loader:
        print(batch_image.shape,batch_label)
